from sqlalchemy import create_engine, inspect, text
from sqlalchemy.exc import SQLAlchemyError
import uuid


class ConnectionManager:
    """管理所有数据库连接（最小可用版）"""

    def __init__(self):
        self.connections = {}          # {conn_id: dict}
        self.engines = {}              # {conn_id: Engine}
        self.current_connection: str | None = None

    # ----------------------------------------------------------
    #  1. 添加连接
    # ----------------------------------------------------------
    def add_connection(self, info: dict) -> str:
        """
        info 字段：
        {
            "name": "本地MySQL",
            "type": "mysql",           # mysql / postgresql / sqlite
            "host": "localhost",
            "port": 3306,
            "user": "root",
            "password": "123456",
            "database": "test"
        }
        """
        conn_id = str(uuid.uuid4())
        self.connections[conn_id] = info
        return conn_id

    # ----------------------------------------------------------
    #  2. 获取 SQLAlchemy Engine（按需创建）
    # ----------------------------------------------------------
    def get_engine(self, conn_id: str):
        if conn_id in self.engines:
            return self.engines[conn_id]

        info = self.connections[conn_id]
        db_type = info["type"].lower()

        if db_type == "mysql":
            url = f"mysql+pymysql://{info['user']}:{info['password']}@{info['host']}:{info['port']}/{info['database']}"
        elif db_type == "postgresql":
            url = f"postgresql+psycopg2://{info['user']}:{info['password']}@{info['host']}:{info['port']}/{info['database']}"
        elif db_type == "sqlite":
            url = f"sqlite:///{info.get('database', ':memory:')}"
        else:
            raise ValueError(f"不支持的数据库类型: {db_type}")

        engine = create_engine(url, pool_pre_ping=True)
        self.engines[conn_id] = engine
        return engine

    # ----------------------------------------------------------
    #  3. 执行任意 SQL（通用入口）
    # ----------------------------------------------------------
    def execute_query(self, sql: str, conn_id: str | None = None):
        """
        统一执行 SELECT / INSERT / UPDATE / DELETE
        返回 (rows: list[dict], message: str)
        """
        if conn_id is None:
            conn_id = self.current_connection
        if conn_id is None:
            raise RuntimeError("未选择数据库连接")

        engine = self.get_engine(conn_id)
        try:
            with engine.connect() as conn:
                result = conn.execute(text(sql))
                if result.returns_rows:
                    rows = [dict(row._mapping) for row in result.fetchall()]
                    return rows, f"查询成功，返回 {len(rows)} 行"
                else:
                    conn.commit()
                    return [], f"执行成功，影响 {result.rowcount} 行"
        except Exception as e:
            return [], str(e)

    # ----------------------------------------------------------
    #  4. 断开连接（可选）
    # ----------------------------------------------------------
    def disconnect(self, conn_id: str):
        if conn_id in self.engines:
            self.engines[conn_id].dispose()
            del self.engines[conn_id]